import json
import os
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor, Future
from typing import Any, Dict, List, Optional, Tuple, Union

from m3u8_downloader import M3U8Downloader


class _DefaultPrintHandler:
    def on_download_start(self, download_id: str = "", **kwargs):
        print(f"[start] id={download_id}")

    def on_playlist_resolved(self, playlist_url: str, download_id: str = "", **kwargs):
        print(f"[resolved] id={download_id} url={playlist_url}")

    def on_key_parsed(self, method: str, uri: str, iv: Optional[str], download_id: str = "", **kwargs):
        print(f"[key] id={download_id} method={method} uri={uri} iv={iv}")

    def on_segment_downloaded(self, index: int, total: int, bytes_downloaded: int, speed_bps: float,
                              download_id: str = "", **kwargs):
        print(f"[seg] id={download_id} {index+1}/{total} bytes={bytes_downloaded} speed={speed_bps:.2f}B/s")

    def on_merge_progress(self, completed: int, total: int, download_id: str = "", **kwargs):
        print(f"[merge] id={download_id} {completed}/{total}")

    def on_completed(self, output: str, download_id: str = "", **kwargs):
        print(f"[done] id={download_id} output={output}")

    def on_error(self, error: Exception, download_id: str = "", **kwargs):
        print(f"[error] id={download_id} {error}")


class DownloadManager:
    """
    管理多个 M3U8Downloader 实例：
    - 每次接收单个 JSON（仅支持单链接）创建一个独立下载任务；
    - 使用 UUID 字符串作为任务 ID（`download_id`），并在 `downloads/<UUID>` 子目录下保存该任务的文件；
    - 支持异步提交与状态查询，并提供暂停/恢复控制接口。

    单任务 JSON 字段示例：
    {
      "url": "...m3u8",
      "output_dir": "downloads",
      "download_id": "<可选指定UUID字符串>",
      "workers": 6,
      "timeout": 10,
      "retries": 3,
      "proxy": null,
      "headers": {"User-Agent": "..."},
      "live_seconds": null,
      "live_limit": null,
      "merge_mp4": true,
      "mp4_name": "merged.mp4",
      "cleanup_ts": true
    }

    说明：当前不支持批量 `urls` 字段，若提供将抛出错误。
    """

    def __init__(self, max_concurrent: int = 2, default_handler: Optional[object] = None):
        """初始化下载管理器。

        - `max_concurrent`：并发运行的任务上限；
        - `default_handler`：默认事件处理器，不传则使用 `_DefaultPrintHandler`。
        """
        self._executor = ThreadPoolExecutor(max_workers=max_concurrent)
        self._default_handler = default_handler or _DefaultPrintHandler()
        self._futures: Dict[str, Future] = {}
        self._results: Dict[str, Dict[str, Any]] = {}
        self._lock = threading.Lock()
        self._downloaders: Dict[str, M3U8Downloader] = {}
        
    def _alloc_id(self) -> str:
        """生成一个 UUID4 字符串作为下载任务 ID。"""
        return str(uuid.uuid4())

    def _parse_json(self, param_json: str) -> Dict[str, Any]:
        cfg = json.loads(param_json)
        if not isinstance(cfg, dict):
            raise ValueError("JSON 根应为对象")
        return cfg

    def _ensure_output_dir(self, path: str) -> str:
        """确保输出目录存在并返回该路径。"""
        os.makedirs(path, exist_ok=True)
        return path

    def _build_downloader(self, url: str, cfg: Dict[str, Any], handler: Optional[object]) -> Tuple[M3U8Downloader, Dict[str, Any], str]:
        """根据配置构建下载器实例与运行参数，并分配/确定 UUID 任务 ID。

        返回三元组：(downloader, run_args, download_id)。
        """
        # 下载器参数
        workers = int(cfg.get("workers", 8))
        timeout = int(cfg.get("timeout", 10))
        retries = int(cfg.get("retries", 3))
        proxy = cfg.get("proxy")
        headers = cfg.get("headers")
        # 拉流轮询周期（直播）
        live_poll_interval = int(cfg.get("poll", cfg.get("live_poll_interval", 5)))
        # 直播/分片限制（传给 download 阶段）
        live_seconds = cfg.get("live_seconds")
        live_limit = cfg.get("live_limit")

        # 运行参数
        merge_mp4 = bool(cfg.get("merge_mp4", False))
        mp4_name = cfg.get("mp4_name")
        cleanup_ts = bool(cfg.get("cleanup_ts", False))

        # 下载 ID 分配
        explicit_id = cfg.get("download_id")
        download_id = str(explicit_id) if explicit_id is not None else self._alloc_id()

        # 通用配置（按 UUID 子目录保存）
        base_output_dir = cfg.get("output_dir") or os.path.join(os.getcwd(), "downloads")
        output_dir = os.path.join(base_output_dir, download_id)
        self._ensure_output_dir(output_dir)

        downloader = M3U8Downloader(
            url=url,
            output_dir=output_dir,
            max_workers=workers,
            timeout=timeout,
            retries=retries,
            headers=headers,
            proxy=proxy,
            live_poll_interval=live_poll_interval,
            download_id=download_id,
            event_handler=(handler or self._default_handler),
        )

        run_args = dict(
            live_duration_seconds=live_seconds,
            segment_limit=live_limit,
            merge_mp4=merge_mp4,
            mp4_name=mp4_name,
            cleanup_ts=cleanup_ts,
        )
        return downloader, run_args, download_id

    def _run_task(self, downloader: M3U8Downloader, run_args: Dict[str, Any]) -> Dict[str, Any]:
        """实际运行下载任务，返回结果字典。"""
        output = downloader.download(**run_args)
        return {"output": output}

    def submit_json(self, param_json: str, handler: Optional[object] = None) -> str:
        """
        提交单任务 JSON：返回 UUID 字符串 `download_id`。

        任务异步执行，可使用 `get_status`/`result`/`list_tasks` 查询。
        若 `cfg` 中提供 `download_id` 则沿用；否则自动生成。
        """
        cfg = self._parse_json(param_json)

        if "urls" in cfg:
            raise ValueError("不支持批量：请仅提供单个 url 字段")
        if "url" not in cfg or not cfg["url"]:
            raise ValueError("缺少 url")

        url = str(cfg["url"]) 
        downloader, run_args, download_id = self._build_downloader(url, cfg, handler)

        fut = self._executor.submit(self._run_task, downloader, run_args)
        with self._lock:
            self._futures[download_id] = fut
            self._downloaders[download_id] = downloader

        def _cb(f: Future, did: str = download_id):
            try:
                res = f.result()
            except Exception as e:
                res = {"error": e}
            with self._lock:
                self._results[did] = res

        fut.add_done_callback(_cb)
        return download_id

    def get_status(self, download_id: str) -> str:
        """查询指定任务的状态：`unknown`/`pending`/`running`/`done`。"""
        with self._lock:
            fut = self._futures.get(download_id)
            if not fut:
                return "unknown"
            if fut.running():
                return "running"
            if fut.done():
                return "done"
            return "pending"

    def result(self, download_id: str) -> Optional[Dict[str, Any]]:
        """获取指定任务结果（若已完成或出错），否则返回 None。"""
        with self._lock:
            return self._results.get(download_id)

    def list_tasks(self) -> List[Dict[str, Any]]:
        """列出当前管理器中所有任务的状态与（可能的）结果。"""
        with self._lock:
            out = []
            for did, fut in self._futures.items():
                state = "running" if fut.running() else ("done" if fut.done() else "pending")
                out.append({"download_id": did, "state": state, "result": self._results.get(did)})
            return out

    def shutdown(self, wait: bool = True):
        """关闭内部线程池。"""
        self._executor.shutdown(wait=wait)

    # ------------------------ 控制接口 ------------------------
    def pause(self, download_id: str) -> bool:
        """暂停指定任务：成功返回 True，找不到或失败返回 False。"""
        with self._lock:
            d = self._downloaders.get(download_id)
        if not d:
            return False
        try:
            d.pause()
            return True
        except Exception:
            return False

    def resume(self, download_id: str) -> bool:
        """恢复指定任务：成功返回 True，找不到或失败返回 False。"""
        with self._lock:
            d = self._downloaders.get(download_id)
        if not d:
            return False
        try:
            d.resume()
            return True
        except Exception:
            return False

    def stop(self, download_id: str, merge_mp4: bool = False) -> bool:
        """停止指定任务：协作式停止；可选择停止后合并已下载为 MP4。

        返回 True 表示成功发出停止请求；找不到或失败返回 False。
        """
        with self._lock:
            d = self._downloaders.get(download_id)
        if not d:
            return False
        try:
            # 向下载器发出停止请求，由其在安全点退出并可选合并 MP4
            d.stop(merge_mp4=merge_mp4)
            return True
        except Exception:
            return False